// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/array_nested.h>

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <nebula/array/array_base.h>
#include <nebula/array/array_primitive.h>
#include <nebula/array/builder_base.h>
#include <nebula/array/builder_nested.h>
#include <nebula/array/concatenate.h>
#include <nebula/array/util.h>
#include <nebula/core/buffer.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_fwd.h>
#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <nebula/bits/bitmap_generate.h>
#include <nebula/bits/bitmap_ops.h>
#include <turbo/base/checked_cast.h>
#include <nebula/util/list_util.h>
#include <turbo/log/logging.h>
#include <turbo/log/logging.h>

namespace nebula {

    using internal::BitmapAnd;
    using internal::CopyBitmap;

    // ----------------------------------------------------------------------
    // ListArray / LargeListArray / ListViewArray / LargeListViewArray (common utilities)

    namespace {

        /// \brief Clean offsets when their null_count is greater than 0
        ///
        /// \pre offsets.null_count() > 0
        template<typename TYPE>
        turbo::Result<BufferVector> CleanListOffsets(const std::shared_ptr<Buffer> &validity_buffer,
                                                     const Array &offsets, MemoryPool *pool) {
            using offset_type = typename TYPE::offset_type;
            using OffsetArrowType = typename CTypeTraits<offset_type>::ArrowType;
            using OffsetArrayType = typename TypeTraits<OffsetArrowType>::ArrayType;

                    DKCHECK_GT(offsets.null_count(), 0);
            const int64_t num_offsets = offsets.length();

            if (!offsets.is_valid(num_offsets - 1)) {
                return turbo::invalid_argument_error("Last list offset should be non-null");
            }

            TURBO_MOVE_OR_RAISE(auto clean_offsets,
                                allocate_buffer(num_offsets * sizeof(offset_type), pool));

            // Copy valid bits, ignoring the final offset (since for a length N list array,
            // we have N + 1 offsets)
            TURBO_MOVE_OR_RAISE(
                    auto clean_validity_buffer,
                    CopyBitmap(pool, offsets.null_bitmap()->data(), offsets.offset(), num_offsets - 1));

            const offset_type *raw_offsets =
                    turbo::checked_cast<const OffsetArrayType &>(offsets).raw_values();
            auto clean_raw_offsets = reinterpret_cast<offset_type *>(clean_offsets->mutable_data());

            // Must work backwards so we can tell how many values were in the last non-null value
            offset_type current_offset = raw_offsets[num_offsets - 1];
            for (int64_t i = num_offsets - 1; i >= 0; --i) {
                if (offsets.is_valid(i)) {
                    current_offset = raw_offsets[i];
                }
                clean_raw_offsets[i] = current_offset;
            }

            return BufferVector({std::move(clean_validity_buffer), std::move(clean_offsets)});
        }

        template<typename TYPE>
        turbo::Result<std::shared_ptr<typename TypeTraits<TYPE>::ArrayType>> ListArrayFromArrays(
                std::shared_ptr<DataType> type, const Array &offsets, const Array &values,
                MemoryPool *pool, std::shared_ptr<Buffer> null_bitmap = nullptr,
                int64_t null_count = kUnknownNullCount) {
            using offset_type = typename TYPE::offset_type;
            using ArrayType = typename TypeTraits<TYPE>::ArrayType;
            using OffsetArrowType = typename CTypeTraits<offset_type>::ArrowType;

            if (offsets.length() == 0) {
                return turbo::invalid_argument_error("List offsets must have non-zero length");
            }

            if (offsets.type_id() != OffsetArrowType::type_id) {
                return turbo::failed_precondition_error("List offsets must be ", OffsetArrowType::type_name());
            }

            if (null_bitmap != nullptr && offsets.data()->may_have_nulls()) {
                return turbo::invalid_argument_error(
                        "Ambiguous to specify both validity map and offsets with nulls");
            }

            if (null_bitmap != nullptr && offsets.offset() != 0) {
                return turbo::unimplemented_error("Null bitmap with offsets slice not supported.");
            }

            // Clean the offsets if they contain nulls.
            if (offsets.null_count() > 0) {
                TURBO_MOVE_OR_RAISE(auto buffers,
                                    CleanListOffsets<TYPE>(null_bitmap, offsets, pool));
                auto data = ArrayData::create(type, offsets.length() - 1, std::move(buffers),
                                              {values.data()}, offsets.null_count(), /*offset=*/0);
                return std::make_shared<ArrayType>(std::move(data));
            }

            using OffsetArrayType = typename TypeTraits<OffsetArrowType>::ArrayType;
            const auto &typed_offsets = turbo::checked_cast<const OffsetArrayType &>(offsets);
            auto buffers = BufferVector({std::move(null_bitmap), typed_offsets.values()});
            auto data = ArrayData::create(type, offsets.length() - 1, std::move(buffers),
                                          {values.data()}, null_count, offsets.offset());
            return std::make_shared<ArrayType>(std::move(data));
        }

        template<typename TYPE>
        turbo::Result<std::shared_ptr<typename TypeTraits<TYPE>::ArrayType>> ListViewArrayFromArrays(
                std::shared_ptr<DataType> type, const Array &offsets, const Array &sizes,
                const Array &values, MemoryPool *pool, std::shared_ptr<Buffer> null_bitmap = nullptr,
                int64_t null_count = kUnknownNullCount) {
            using offset_type = typename TYPE::offset_type;
            using ArrayType = typename TypeTraits<TYPE>::ArrayType;
            using OffsetArrowType = typename CTypeTraits<offset_type>::ArrowType;

            if (offsets.type_id() != OffsetArrowType::type_id) {
                return turbo::failed_precondition_error("List offsets must be ", OffsetArrowType::type_name());
            }

            if (sizes.length() != offsets.length() && sizes.length() != offsets.length() - 1) {
                return turbo::invalid_argument_error(
                        "List sizes must have the same length as offsets or one less than offsets");
            }
            if (sizes.type_id() != OffsetArrowType::type_id) {
                return turbo::failed_precondition_error("List sizes must be ", OffsetArrowType::type_name());
            }

            if (offsets.offset() != sizes.offset()) {
                return turbo::invalid_argument_error("List offsets and sizes must have the same offset");
            }
            const int64_t array_offset = sizes.offset();

            if (null_bitmap) {
                if (offsets.null_count() > 0 || sizes.null_count() > 0) {
                    return turbo::invalid_argument_error(
                            "Ambiguous to specify both validity map and offsets or sizes with nulls");
                }
                if (array_offset != 0) {
                    return turbo::invalid_argument_error(
                            "List offsets and sizes must not be slices if a validity map is specified");
                }
            } else {
                if (offsets.null_count() > 0 && sizes.null_count() > 0) {
                    return turbo::invalid_argument_error("Ambiguous to specify both offsets and sizes with nulls");
                }
            }

                    DKCHECK(offsets.length() == sizes.length() || offsets.length() - 1 == sizes.length());

            using OffsetArrayType = typename TypeTraits<OffsetArrowType>::ArrayType;
            const auto &typed_offsets = turbo::checked_cast<const OffsetArrayType &>(offsets);
            const auto &typed_sizes = turbo::checked_cast<const OffsetArrayType &>(sizes);

            auto derived_validity_buffer = std::move(null_bitmap);
            if (offsets.null_count() > 0) {
                derived_validity_buffer = offsets.null_bitmap();
                null_count = offsets.null_count();
                // We allow construction from an offsets array containing one extra value.
                // If that is the case, we might need to discount one null from out_null_count.
                if (offsets.length() - 1 == sizes.length() && !offsets.is_valid(sizes.length())) {
                    null_count -= 1;
                }
            } else if (sizes.null_count() > 0) {
                derived_validity_buffer = sizes.null_bitmap();
                null_count = sizes.null_count();
            }

            auto buffers = BufferVector({
                                                std::move(derived_validity_buffer),
                                                typed_offsets.values(),
                                                typed_sizes.values(),
                                        });
            auto data = ArrayData::create(type, sizes.length(), std::move(buffers), {values.data()},
                                          null_count, array_offset);
            return std::make_shared<ArrayType>(std::move(data));
        }

        static std::shared_ptr<Array> SliceArrayWithOffsets(const Array &array, int64_t begin,
                                                            int64_t end) {
            return array.slice(begin, end - begin);
        }

        template<typename ListArrayT>
        turbo::Result<std::shared_ptr<Array>> FlattenListArray(const ListArrayT &list_array,
                                                               MemoryPool *memory_pool) {
            const int64_t list_array_length = list_array.length();
            std::shared_ptr<nebula::Array> value_array = list_array.values();

            // Shortcut: if a ListArray does not contain nulls, then simply slice its
            // value array with the first and the last offsets.
            if (list_array.null_count() == 0) {
                return SliceArrayWithOffsets(*value_array, list_array.value_offset(0),
                                             list_array.value_offset(list_array_length));
            }

            // Second shortcut: if the list array is *all* nulls, then just return
            // an empty array.
            if (list_array.null_count() == list_array.length()) {
                return MakeEmptyArray(value_array->type(), memory_pool);
            }

            // The ListArray contains nulls: there may be a non-empty sub-list behind
            // a null and it must not be contained in the result.
            std::vector<std::shared_ptr<Array>> non_null_fragments;
            int64_t valid_begin = 0;
            while (valid_begin < list_array_length) {
                int64_t valid_end = valid_begin;
                while (valid_end < list_array_length &&
                       (list_array.is_valid(valid_end) || list_array.value_length(valid_end) == 0)) {
                    ++valid_end;
                }
                if (valid_begin < valid_end) {
                    non_null_fragments.push_back(
                            SliceArrayWithOffsets(*value_array, list_array.value_offset(valid_begin),
                                                  list_array.value_offset(valid_end)));
                }
                valid_begin = valid_end + 1;  // skip null entry
            }

            // Final attempt to avoid invoking concatenate().
            if (non_null_fragments.size() == 1) {
                return non_null_fragments[0];
            } else if (non_null_fragments.size() == 0) {
                return MakeEmptyArray(value_array->type(), memory_pool);
            }

            return concatenate(non_null_fragments, memory_pool);
        }

        template<typename ListViewArrayT, bool HasNulls>
        turbo::Result<std::shared_ptr<Array>> FlattenListViewArray(const ListViewArrayT &list_view_array,
                                                                   MemoryPool *memory_pool) {
            using offset_type = typename ListViewArrayT::offset_type;
            const int64_t list_view_array_offset = list_view_array.offset();
            const int64_t list_view_array_length = list_view_array.length();
            std::shared_ptr<nebula::Array> value_array = list_view_array.values();

            if (list_view_array_length == 0) {
                return SliceArrayWithOffsets(*value_array, 0, 0);
            }

            // If the list array is *all* nulls, then just return an empty array.
            if constexpr (HasNulls) {
                if (list_view_array.null_count() == list_view_array.length()) {
                    return MakeEmptyArray(value_array->type(), memory_pool);
                }
            }

            const auto *validity = list_view_array.data()->template get_values<uint8_t>(0, 0);
            const auto *offsets = list_view_array.data()->template get_values<offset_type>(1);
            const auto *sizes = list_view_array.data()->template get_values<offset_type>(2);

            auto is_null_or_empty = [&](int64_t i) {
                if (HasNulls && !bit_util::get_bit(validity, list_view_array_offset + i)) {
                    return true;
                }
                return sizes[i] == 0;
            };

            // Index of the first valid, non-empty list-view.
            int64_t first_i = 0;
            for (; first_i < list_view_array_length; first_i++) {
                if (!is_null_or_empty(first_i)) {
                    break;
                }
            }
            // If all list-views are empty, return an empty array.
            if (first_i == list_view_array_length) {
                return MakeEmptyArray(value_array->type(), memory_pool);
            }

            std::vector<std::shared_ptr<Array>> slices;
            {
                int64_t i = first_i;
                auto begin_offset = offsets[i];
                auto end_offset = offsets[i] + sizes[i];
                i += 1;
                // Inductive invariant: slices and the always non-empty values slice
                // [begin_offset, end_offset) contains all the maximally contiguous slices of the
                // values array that are covered by all the list-views before list-view i.
                for (; i < list_view_array_length; i++) {
                    if (is_null_or_empty(i)) {
                        // The invariant is preserved by simply preserving the current set of slices.
                    } else {
                        if (offsets[i] == end_offset) {
                            end_offset += sizes[i];
                            // The invariant is preserved because since the non-empty list-view i
                            // starts at end_offset, the current range can be extended to end at
                            // offsets[i] + sizes[i] (the same as end_offset + sizes[i]).
                        } else {
                            // The current slice can't be extended because the list-view i either
                            // shares values with the current slice or starts after the position
                            // immediately after the end of the current slice.
                            slices.push_back(SliceArrayWithOffsets(*value_array, begin_offset, end_offset));
                            begin_offset = offsets[i];
                            end_offset = offsets[i] + sizes[i];
                            // The invariant is preserved because a maximally contiguous slice of
                            // the values array (i.e. one that can't be extended) was added to slices
                            // and [begin_offset, end_offset) is non-empty and contains the
                            // current list-view i.
                        }
                    }
                }
                slices.push_back(SliceArrayWithOffsets(*value_array, begin_offset, end_offset));
            }

            // Final attempt to avoid invoking concatenate().
            switch (slices.size()) {
                case 0:
                    return MakeEmptyArray(value_array->type(), memory_pool);
                case 1:
                    return slices[0];
            }

            return concatenate(slices, memory_pool);
        }

        std::shared_ptr<Array> BoxOffsets(const std::shared_ptr<DataType> &boxed_type,
                                          const ArrayData &data) {
            const int64_t num_offsets =
                    is_list_view(data.type->id()) ? data.length : data.length + 1;
            std::vector<std::shared_ptr<Buffer>> buffers = {nullptr, data.buffers[1]};
            auto offsets_data =
                    std::make_shared<ArrayData>(boxed_type, /*length=*/num_offsets, std::move(buffers),
                            /*null_count=*/0, data.offset);
            return make_array(offsets_data);
        }

        std::shared_ptr<Array> BoxSizes(const std::shared_ptr<DataType> &boxed_type,
                                        const ArrayData &data) {
                    DKCHECK(is_list_view(data.type->id()));
            std::vector<std::shared_ptr<Buffer>> buffers = {nullptr, data.buffers[2]};
            auto sizes_data =
                    std::make_shared<ArrayData>(boxed_type, data.length, std::move(buffers),
                            /*null_count=*/0, data.offset);
            return make_array(sizes_data);
        }

        template<typename DestListViewType, typename SrcListType>
        turbo::Result<std::shared_ptr<ArrayData>> ListViewFromListImpl(
                const std::shared_ptr<ArrayData> &list_data, MemoryPool *pool) {
            static_assert(
                    std::is_same<typename SrcListType::offset_type,
                            typename DestListViewType::offset_type>::value,
                    "Offset types between list type and list-view type are expected to match");
            using offset_type = typename SrcListType::offset_type;
            const auto &list_type = turbo::checked_cast<const SrcListType &>(*list_data->type);

            // To re-use the validity and offsets buffers, a sizes buffer with enough
            // padding on the beginning is allocated and filled with the sizes after
            // list_data->offset.
            const int64_t buffer_length = list_data->offset + list_data->length;
            TURBO_MOVE_OR_RAISE(auto sizes_buffer,
                                allocate_buffer(buffer_length * sizeof(offset_type), pool));
            const auto *offsets = list_data->template get_values<offset_type>(1, 0);
            auto *sizes = sizes_buffer->mutable_data_as<offset_type>();
            // Zero the initial padding area to avoid leaking any data when buffers are
            // sent over IPC or through the C Data interface.
            memset(sizes, 0, list_data->offset * sizeof(offset_type));
            for (int64_t i = list_data->offset; i < buffer_length; i++) {
                sizes[i] = offsets[i + 1] - offsets[i];
            }
            BufferVector buffers = {list_data->buffers[0], list_data->buffers[1],
                                    std::move(sizes_buffer)};

            return ArrayData::create(std::make_shared<DestListViewType>(list_type.get_value_type()),
                                     list_data->length, std::move(buffers),
                                     {list_data->child_data[0]}, list_data->null_count,
                                     list_data->offset);
        }

        template<typename DestListType, typename SrcListViewType>
        turbo::Result<std::shared_ptr<ArrayData>> ListFromListViewImpl(
                const std::shared_ptr<ArrayData> &list_view_data, MemoryPool *pool) {
            static_assert(
                    std::is_same<typename SrcListViewType::offset_type,
                            typename DestListType::offset_type>::value,
                    "Offset types between list type and list-view type are expected to match");
            using offset_type = typename DestListType::offset_type;
            using ListBuilderType = typename TypeTraits<DestListType>::BuilderType;

            const auto &list_view_type =
                    turbo::checked_cast<const SrcListViewType &>(*list_view_data->type);
            const auto &value_type = list_view_type.get_value_type();
            const auto list_type = std::make_shared<DestListType>(value_type);

            TURBO_MOVE_OR_RAISE(auto sum_of_list_view_sizes,
                                list_util::internal::SumOfLogicalListSizes(*list_view_data));
            TURBO_MOVE_OR_RAISE(std::shared_ptr<ArrayBuilder> value_builder,
                                MakeBuilder(value_type, pool));
            TURBO_RETURN_NOT_OK(value_builder->Reserve(sum_of_list_view_sizes));
            auto list_builder = std::make_shared<ListBuilderType>(pool, value_builder, list_type);
            TURBO_RETURN_NOT_OK(list_builder->Reserve(list_view_data->length));

            ArraySpan values{*list_view_data->child_data[0]};
            const auto *in_validity_bitmap = list_view_data->get_values<uint8_t>(0);
            const auto *in_offsets = list_view_data->get_values<offset_type>(1);
            const auto *in_sizes = list_view_data->get_values<offset_type>(2);
            for (int64_t i = 0; i < list_view_data->length; ++i) {
                const bool is_valid =
                        !in_validity_bitmap ||
                        bit_util::get_bit(in_validity_bitmap, list_view_data->offset + i);
                const int64_t size = is_valid ? in_sizes[i] : 0;
                TURBO_RETURN_NOT_OK(list_builder->append(is_valid, size));
                TURBO_RETURN_NOT_OK(value_builder->append_array_slice(values, in_offsets[i], size));
            }
            std::shared_ptr<ArrayData> list_array_data;
            TURBO_RETURN_NOT_OK(list_builder->finish_internal(&list_array_data));
            return list_array_data;
        }

    }  // namespace

    namespace internal {

        template<typename TYPE>
        inline void SetListData(VarLengthListLikeArray<TYPE> *self,
                                const std::shared_ptr<ArrayData> &data,
                                Type::type expected_type_id) {
            KCHECK_EQ(data->buffers.size(), is_list_view(TYPE::type_id) ? 3 : 2);
            KCHECK_EQ(data->type->id(), expected_type_id);
            KCHECK_EQ(data->child_data.size(), 1);

            self->Array::set_data(data);

            self->list_type_ = turbo::checked_cast<const TYPE *>(data->type.get());
            self->raw_value_offsets_ =
                    data->GetValuesSafe<typename TYPE::offset_type>(1, /*offset=*/0);
            // BaseListViewArray::set_data takes care of setting raw_value_sizes_.

            KCHECK_EQ(self->list_type_->get_value_type()->id(), data->child_data[0]->type->id());
                    DKCHECK(self->list_type_->get_value_type()->equals(data->child_data[0]->type));
            self->values_ = make_array(self->data_->child_data[0]);
        }

        turbo::Result<std::shared_ptr<Array>> FlattenLogicalListRecursively(const Array &in_array,
                                                                            MemoryPool *memory_pool) {
            std::shared_ptr<Array> array = in_array.slice(0, in_array.length());
            for (auto kind = array->type_id(); is_list(kind) || is_list_view(kind);
                 kind = array->type_id()) {
                switch (kind) {
                    case Type::LIST: {
                        TURBO_MOVE_OR_RAISE(
                                array, (turbo::checked_cast<const ListArray *>(array.get())->Flatten(memory_pool)));
                        break;
                    }
                    case Type::LARGE_LIST: {
                        TURBO_MOVE_OR_RAISE(
                                array,
                                (turbo::checked_cast<const LargeListArray *>(array.get())->Flatten(memory_pool)));
                        break;
                    }
                    case Type::LIST_VIEW: {
                        TURBO_MOVE_OR_RAISE(
                                array,
                                (turbo::checked_cast<const ListViewArray *>(array.get())->Flatten(memory_pool)));
                        break;
                    }
                    case Type::LARGE_LIST_VIEW: {
                        TURBO_MOVE_OR_RAISE(
                                array,
                                (turbo::checked_cast<const LargeListViewArray *>(array.get())->Flatten(memory_pool)));
                        break;
                    }
                    case Type::FIXED_SIZE_LIST: {
                        TURBO_MOVE_OR_RAISE(
                                array,
                                (turbo::checked_cast<const FixedSizeListArray *>(array.get())->Flatten(memory_pool)));
                        break;
                    }
                    default:
                        unreachable("unexpected non-list type");
                        break;
                }
            }
            return array;
        }

    }  // namespace internal

    // ----------------------------------------------------------------------
    // ListArray

    ListArray::ListArray(std::shared_ptr<ArrayData> data) {
        ListArray::set_data(std::move(data));
    }

    ListArray::ListArray(std::shared_ptr<DataType> type, int64_t length,
                         std::shared_ptr<Buffer> value_offsets, std::shared_ptr<Array> values,
                         std::shared_ptr<Buffer> null_bitmap, int64_t null_count,
                         int64_t offset) {
        KCHECK_EQ(type->id(), Type::LIST);
        auto internal_data = ArrayData::create(
                std::move(type), length,
                BufferVector{std::move(null_bitmap), std::move(value_offsets)}, null_count, offset);
        internal_data->child_data.emplace_back(values->data());
        set_data(std::move(internal_data));
    }

    void ListArray::set_data(const std::shared_ptr<ArrayData> &data) {
        internal::SetListData(this, data);
    }

    turbo::Result<std::shared_ptr<ListArray>> ListArray::from_arrays(
            const Array &offsets, const Array &values, MemoryPool *pool,
            std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        return ListArrayFromArrays<ListType>(std::make_shared<ListType>(values.type()), offsets,
                                             values, pool, null_bitmap, null_count);
    }

    turbo::Result<std::shared_ptr<ListArray>> ListArray::FromListView(const ListViewArray &source,
                                                                      MemoryPool *pool) {
        TURBO_MOVE_OR_RAISE(
                auto data, (ListFromListViewImpl<ListType, ListViewType>(source.data(), pool)));
        return std::make_shared<ListArray>(std::move(data));
    }

    turbo::Result<std::shared_ptr<ListArray>> ListArray::from_arrays(
            std::shared_ptr<DataType> type, const Array &offsets, const Array &values,
            MemoryPool *pool, std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        if (type->id() != Type::LIST) {
            return turbo::failed_precondition_error("Expected list type, got ", type->to_string());
        }
        const auto &list_type = turbo::checked_cast<const ListType &>(*type);
        if (!list_type.get_value_type()->equals(values.type())) {
            return turbo::failed_precondition_error("Mismatching list value type");
        }
        return ListArrayFromArrays<ListType>(std::move(type), offsets, values, pool,
                                             null_bitmap, null_count);
    }

    turbo::Result<std::shared_ptr<Array>> ListArray::Flatten(MemoryPool *memory_pool) const {
        return FlattenListArray(*this, memory_pool);
    }

    std::shared_ptr<Array> ListArray::offsets() const { return BoxOffsets(int32(), *data_); }

    // ----------------------------------------------------------------------
    // LargeListArray

    LargeListArray::LargeListArray(const std::shared_ptr<ArrayData> &data) {
        LargeListArray::set_data(data);
    }

    LargeListArray::LargeListArray(const std::shared_ptr<DataType> &type, int64_t length,
                                   const std::shared_ptr<Buffer> &value_offsets,
                                   const std::shared_ptr<Array> &values,
                                   const std::shared_ptr<Buffer> &null_bitmap,
                                   int64_t null_count, int64_t offset) {
        KCHECK_EQ(type->id(), Type::LARGE_LIST);
        auto internal_data =
                ArrayData::create(type, length, {null_bitmap, value_offsets}, null_count, offset);
        internal_data->child_data.emplace_back(values->data());
        LargeListArray::set_data(internal_data);
    }

    void LargeListArray::set_data(const std::shared_ptr<ArrayData> &data) {
        internal::SetListData(this, data);
    }

    turbo::Result<std::shared_ptr<LargeListArray>> LargeListArray::from_arrays(
            const Array &offsets, const Array &values, MemoryPool *pool,
            std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        return ListArrayFromArrays<LargeListType>(
                std::make_shared<LargeListType>(values.type()), offsets, values, pool, null_bitmap,
                null_count);
    }

    turbo::Result<std::shared_ptr<LargeListArray>> LargeListArray::FromListView(
            const LargeListViewArray &source, MemoryPool *pool) {
        TURBO_MOVE_OR_RAISE(
                auto data,
                (ListFromListViewImpl<LargeListType, LargeListViewType>(source.data(), pool)));
        return std::make_shared<LargeListArray>(std::move(data));
    }

    turbo::Result<std::shared_ptr<LargeListArray>> LargeListArray::from_arrays(
            std::shared_ptr<DataType> type, const Array &offsets, const Array &values,
            MemoryPool *pool, std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        if (type->id() != Type::LARGE_LIST) {
            return turbo::failed_precondition_error("Expected large list type, got ", type->to_string());
        }
        const auto &list_type = turbo::checked_cast<const LargeListType &>(*type);
        if (!list_type.get_value_type()->equals(values.type())) {
            return turbo::failed_precondition_error("Mismatching list value type");
        }
        return ListArrayFromArrays<LargeListType>(std::move(type), offsets, values, pool,
                                                  null_bitmap, null_count);
    }

    turbo::Result<std::shared_ptr<Array>> LargeListArray::Flatten(MemoryPool *memory_pool) const {
        return FlattenListArray(*this, memory_pool);
    }

    std::shared_ptr<Array> LargeListArray::offsets() const {
        return BoxOffsets(int64(), *data_);
    }

    // ----------------------------------------------------------------------
    // ListViewArray

    ListViewArray::ListViewArray(std::shared_ptr<ArrayData> data) {
        ListViewArray::set_data(std::move(data));
    }

    ListViewArray::ListViewArray(std::shared_ptr<DataType> type, int64_t length,
                                 std::shared_ptr<Buffer> value_offsets,
                                 std::shared_ptr<Buffer> value_sizes,
                                 std::shared_ptr<Array> values,
                                 std::shared_ptr<Buffer> null_bitmap, int64_t null_count,
                                 int64_t offset) {
        ListViewArray::set_data(ArrayData::create(
                std::move(type), length,
                {std::move(null_bitmap), std::move(value_offsets), std::move(value_sizes)},
                /*child_data=*/{values->data()}, null_count, offset));
    }

    void ListViewArray::set_data(const std::shared_ptr<ArrayData> &data) {
        internal::SetListData(this, data);
        raw_value_sizes_ = data->GetValuesSafe<ListViewType::offset_type>(2, /*offset=*/0);
    }

    turbo::Result<std::shared_ptr<ListViewArray>> ListViewArray::from_arrays(
            const Array &offsets, const Array &sizes, const Array &values, MemoryPool *pool,
            std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        return ListViewArrayFromArrays<ListViewType>(
                std::make_shared<ListViewType>(values.type()), offsets, sizes, values, pool,
                null_bitmap, null_count);
    }

    turbo::Result<std::shared_ptr<ListViewArray>> ListViewArray::from_arrays(
            std::shared_ptr<DataType> type, const Array &offsets, const Array &sizes,
            const Array &values, MemoryPool *pool, std::shared_ptr<Buffer> null_bitmap,
            int64_t null_count) {
        if (type->id() != Type::LIST_VIEW) {
            return turbo::failed_precondition_error("Expected list-view type, got ", type->to_string());
        }
        const auto &list_view_type = turbo::checked_cast<const ListViewType &>(*type);
        if (!list_view_type.get_value_type()->equals(values.type())) {
            return turbo::failed_precondition_error("Mismatching list-view value type");
        }
        return ListViewArrayFromArrays<ListViewType>(std::move(type), offsets, sizes, values,
                                                     pool, null_bitmap, null_count);
    }

    turbo::Result<std::shared_ptr<ListViewArray>> ListViewArray::FromList(const ListArray &source,
                                                                          MemoryPool *pool) {
        TURBO_MOVE_OR_RAISE(
                auto data, (ListViewFromListImpl<ListViewType, ListType>(source.data(), pool)));
        return std::make_shared<ListViewArray>(std::move(data));
    }

    turbo::Result<std::shared_ptr<LargeListViewArray>> LargeListViewArray::FromList(
            const LargeListArray &source, MemoryPool *pool) {
        TURBO_MOVE_OR_RAISE(
                auto data,
                (ListViewFromListImpl<LargeListViewType, LargeListType>(source.data(), pool)));
        return std::make_shared<LargeListViewArray>(std::move(data));
    }

    turbo::Result<std::shared_ptr<Array>> ListViewArray::Flatten(MemoryPool *memory_pool) const {
        if (null_count() > 0) {
            return FlattenListViewArray<ListViewArray, true>(*this, memory_pool);
        }
        return FlattenListViewArray<ListViewArray, false>(*this, memory_pool);
    }

    std::shared_ptr<Array> ListViewArray::offsets() const {
        return BoxOffsets(int32(), *data_);
    }

    std::shared_ptr<Array> ListViewArray::sizes() const { return BoxSizes(int32(), *data_); }

    // ----------------------------------------------------------------------
    // LargeListViewArray

    LargeListViewArray::LargeListViewArray(std::shared_ptr<ArrayData> data) {
        LargeListViewArray::set_data(std::move(data));
    }

    LargeListViewArray::LargeListViewArray(std::shared_ptr<DataType> type, int64_t length,
                                           std::shared_ptr<Buffer> value_offsets,
                                           std::shared_ptr<Buffer> value_sizes,
                                           std::shared_ptr<Array> values,
                                           std::shared_ptr<Buffer> null_bitmap,
                                           int64_t null_count, int64_t offset) {
        LargeListViewArray::set_data(ArrayData::create(
                type, length,
                {std::move(null_bitmap), std::move(value_offsets), std::move(value_sizes)},
                /*child_data=*/{values->data()}, null_count, offset));
    }

    void LargeListViewArray::set_data(const std::shared_ptr<ArrayData> &data) {
        internal::SetListData(this, data);
        raw_value_sizes_ = data->GetValuesSafe<LargeListViewType::offset_type>(2, /*offset=*/0);
    }

    turbo::Result<std::shared_ptr<LargeListViewArray>> LargeListViewArray::from_arrays(
            const Array &offsets, const Array &sizes, const Array &values, MemoryPool *pool,
            std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        return ListViewArrayFromArrays<LargeListViewType>(
                std::make_shared<LargeListViewType>(values.type()), offsets, sizes, values, pool,
                null_bitmap, null_count);
    }

    turbo::Result<std::shared_ptr<LargeListViewArray>> LargeListViewArray::from_arrays(
            std::shared_ptr<DataType> type, const Array &offsets, const Array &sizes,
            const Array &values, MemoryPool *pool, std::shared_ptr<Buffer> null_bitmap,
            int64_t null_count) {
        if (type->id() != Type::LARGE_LIST_VIEW) {
            return turbo::failed_precondition_error("Expected large list-view type, got ", type->to_string());
        }
        const auto &large_list_view_type = turbo::checked_cast<const LargeListViewType &>(*type);
        if (!large_list_view_type.get_value_type()->equals(values.type())) {
            return turbo::failed_precondition_error("Mismatching large list-view value type");
        }
        return ListViewArrayFromArrays<LargeListViewType>(
                std::move(type), offsets, sizes, values, pool, null_bitmap, null_count);
    }

    turbo::Result<std::shared_ptr<Array>> LargeListViewArray::Flatten(
            MemoryPool *memory_pool) const {
        if (null_count() > 0) {
            return FlattenListViewArray<LargeListViewArray, true>(*this, memory_pool);
        }
        return FlattenListViewArray<LargeListViewArray, false>(*this, memory_pool);
    }

    std::shared_ptr<Array> LargeListViewArray::offsets() const {
        return BoxOffsets(int64(), *data_);
    }

    std::shared_ptr<Array> LargeListViewArray::sizes() const {
        return BoxSizes(int64(), *data_);
    }

    // ----------------------------------------------------------------------
    // MapArray

    MapArray::MapArray(const std::shared_ptr<ArrayData> &data) { set_data(data); }

    MapArray::MapArray(const std::shared_ptr<DataType> &type, int64_t length,
                       const std::shared_ptr<Buffer> &offsets,
                       const std::shared_ptr<Array> &values,
                       const std::shared_ptr<Buffer> &null_bitmap, int64_t null_count,
                       int64_t offset) {
        set_data(ArrayData::create(type, length, {null_bitmap, offsets}, {values->data()},
                                  null_count, offset));
    }

    MapArray::MapArray(const std::shared_ptr<DataType> &type, int64_t length,
                       BufferVector buffers, const std::shared_ptr<Array> &keys,
                       const std::shared_ptr<Array> &items, int64_t null_count,
                       int64_t offset) {
        auto pair_data = ArrayData::create(type->fields()[0]->type(), keys->data()->length,
                                           {nullptr}, {keys->data(), items->data()}, 0);
        auto map_data =
                ArrayData::create(type, length, std::move(buffers), {pair_data}, null_count, offset);
        set_data(map_data);
    }

    MapArray::MapArray(const std::shared_ptr<DataType> &type, int64_t length,
                       const std::shared_ptr<Buffer> &offsets,
                       const std::shared_ptr<Array> &keys,
                       const std::shared_ptr<Array> &items,
                       const std::shared_ptr<Buffer> &null_bitmap, int64_t null_count,
                       int64_t offset)
            : MapArray(type, length, {null_bitmap, offsets}, keys, items, null_count, offset) {}

    turbo::Result<std::shared_ptr<Array>> MapArray::FromArraysInternal(
            std::shared_ptr<DataType> type, const std::shared_ptr<Array> &offsets,
            const std::shared_ptr<Array> &keys, const std::shared_ptr<Array> &items,
            MemoryPool *pool, std::shared_ptr<Buffer> null_bitmap) {
        using offset_type = typename MapType::offset_type;
        using OffsetArrowType = typename CTypeTraits<offset_type>::ArrowType;

        if (offsets->length() == 0) {
            return turbo::invalid_argument_error("Map offsets must have non-zero length");
        }

        if (offsets->type_id() != OffsetArrowType::type_id) {
            return turbo::failed_precondition_error("Map offsets must be ", OffsetArrowType::type_name());
        }

        if (keys->null_count() != 0) {
            return turbo::invalid_argument_error("Map cannot contain NULL valued keys");
        }

        if (keys->length() != items->length()) {
            return turbo::invalid_argument_error("Map key and item arrays must be equal length");
        }

        if (null_bitmap != nullptr && offsets->data()->may_have_nulls()) {
            return turbo::invalid_argument_error(
                    "Ambiguous to specify both validity map and offsets with nulls");
        }

        if (null_bitmap != nullptr && offsets->offset() != 0) {
            return turbo::unimplemented_error("Null bitmap with offsets slice not supported.");
        }

        if (offsets->data()->may_have_nulls()) {
            TURBO_MOVE_OR_RAISE(auto buffers,
                                CleanListOffsets<MapType>(nullptr, *offsets, pool));
            return std::make_shared<MapArray>(type, offsets->length() - 1, std::move(buffers),
                                              keys, items, offsets->null_count(), 0);
        }

        using OffsetArrayType = typename TypeTraits<OffsetArrowType>::ArrayType;
        const auto &typed_offsets = turbo::checked_cast<const OffsetArrayType &>(*offsets);

        BufferVector buffers;
        buffers.resize(2);
        int64_t null_count = 0;
        if (null_bitmap) {
            buffers[0] = std::move(null_bitmap);
            null_count = kUnknownNullCount;
        }
        buffers[1] = typed_offsets.values();
        return std::make_shared<MapArray>(type, offsets->length() - 1, std::move(buffers), keys,
                                          items, /*null_count=*/null_count, offsets->offset());
    }

    turbo::Result<std::shared_ptr<Array>> MapArray::from_arrays(const std::shared_ptr<Array> &offsets,
                                                               const std::shared_ptr<Array> &keys,
                                                               const std::shared_ptr<Array> &items,
                                                               MemoryPool *pool,
                                                               std::shared_ptr<Buffer> null_bitmap) {
        return FromArraysInternal(std::make_shared<MapType>(keys->type(), items->type()),
                                  offsets, keys, items, pool, std::move(null_bitmap));
    }

    turbo::Result<std::shared_ptr<Array>> MapArray::from_arrays(std::shared_ptr<DataType> type,
                                                               const std::shared_ptr<Array> &offsets,
                                                               const std::shared_ptr<Array> &keys,
                                                               const std::shared_ptr<Array> &items,
                                                               MemoryPool *pool,
                                                               std::shared_ptr<Buffer> null_bitmap) {
        if (type->id() != Type::MAP) {
            return turbo::failed_precondition_error("Expected map type, got ", type->to_string());
        }
        const auto &map_type = turbo::checked_cast<const MapType &>(*type);
        if (!map_type.key_type()->equals(keys->type())) {
            return turbo::failed_precondition_error("Mismatching map keys type");
        }
        if (!map_type.item_type()->equals(items->type())) {
            return turbo::failed_precondition_error("Mismatching map items type");
        }
        return FromArraysInternal(std::move(type), offsets, keys, items, pool,
                                  std::move(null_bitmap));
    }

    turbo::Status MapArray::ValidateChildData(
            const std::vector<std::shared_ptr<ArrayData>> &child_data) {
        if (child_data.size() != 1) {
            return turbo::invalid_argument_error("Expected one child array for map array");
        }
        const auto &pair_data = child_data[0];
        if (pair_data->type->id() != Type::STRUCT) {
            return turbo::invalid_argument_error("Map array child array should have struct type");
        }
        if (pair_data->may_have_nulls()) {
            return turbo::invalid_argument_error("Map array child array should have no nulls");
        }
        if (pair_data->child_data.size() != 2) {
            return turbo::invalid_argument_error("Map array child array should have two fields");
        }
        if (pair_data->child_data[0]->may_have_nulls()) {
            return turbo::invalid_argument_error("Map array keys array should have no nulls");
        }
        return turbo::OkStatus();
    }

    void MapArray::set_data(const std::shared_ptr<ArrayData> &data) {
        KCHECK_OK(ValidateChildData(data->child_data));

        internal::SetListData(this, data, Type::MAP);
        map_type_ = turbo::checked_cast<const MapType *>(data->type.get());
        const auto &pair_data = data->child_data[0];
        keys_ = make_array(pair_data->child_data[0]);
        items_ = make_array(pair_data->child_data[1]);
    }

    // ----------------------------------------------------------------------
    // FixedSizeListArray

    FixedSizeListArray::FixedSizeListArray(const std::shared_ptr<ArrayData> &data) {
        set_data(data);
    }

    FixedSizeListArray::FixedSizeListArray(const std::shared_ptr<DataType> &type,
                                           int64_t length,
                                           const std::shared_ptr<Array> &values,
                                           const std::shared_ptr<Buffer> &null_bitmap,
                                           int64_t null_count, int64_t offset) {
        auto internal_data = ArrayData::create(type, length, {null_bitmap}, null_count, offset);
        internal_data->child_data.emplace_back(values->data());
        set_data(internal_data);
    }

    void FixedSizeListArray::set_data(const std::shared_ptr<ArrayData> &data) {
        KCHECK_EQ(data->type->id(), Type::FIXED_SIZE_LIST);
        this->Array::set_data(data);

        KCHECK_EQ(list_type()->get_value_type()->id(), data->child_data[0]->type->id());
                DKCHECK(list_type()->get_value_type()->equals(data->child_data[0]->type));
        list_size_ = list_type()->list_size();

        KCHECK_EQ(data_->child_data.size(), 1);
        values_ = make_array(data_->child_data[0]);
    }

    const FixedSizeListType *FixedSizeListArray::list_type() const {
        return turbo::checked_cast<const FixedSizeListType *>(data_->type.get());
    }

    const std::shared_ptr<DataType> &FixedSizeListArray::get_value_type() const {
        return list_type()->get_value_type();
    }

    const std::shared_ptr<Array> &FixedSizeListArray::values() const { return values_; }

    turbo::Result<std::shared_ptr<Array>> FixedSizeListArray::from_arrays(
            const std::shared_ptr<Array> &values, int32_t list_size,
            std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        if (list_size <= 0) {
            return turbo::invalid_argument_error("list_size needs to be a strict positive integer");
        }

        if ((values->length() % list_size) != 0) {
            return turbo::invalid_argument_error(
                    "The length of the values Array needs to be a multiple of the list_size");
        }
        int64_t length = values->length() / list_size;
        auto list_type = std::make_shared<FixedSizeListType>(values->type(), list_size);

        return std::make_shared<FixedSizeListArray>(list_type, length, values, null_bitmap,
                                                    null_count);
    }

    turbo::Result<std::shared_ptr<Array>> FixedSizeListArray::from_arrays(
            const std::shared_ptr<Array> &values, std::shared_ptr<DataType> type,
            std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
        if (type->id() != Type::FIXED_SIZE_LIST) {
            return turbo::failed_precondition_error("Expected fixed size list type, got ", type->to_string());
        }
        const auto &list_type = turbo::checked_cast<const FixedSizeListType &>(*type);

        if (!list_type.get_value_type()->equals(values->type())) {
            return turbo::failed_precondition_error("Mismatching list value type");
        }
        if ((values->length() % list_type.list_size()) != 0) {
            return turbo::invalid_argument_error(
                    "The length of the values Array needs to be a multiple of the list size");
        }
        int64_t length = values->length() / list_type.list_size();

        return std::make_shared<FixedSizeListArray>(type, length, values, null_bitmap,
                                                    null_count);
    }

    turbo::Result<std::shared_ptr<Array>> FixedSizeListArray::Flatten(
            MemoryPool *memory_pool) const {
        return FlattenListArray(*this, memory_pool);
    }

    // ----------------------------------------------------------------------
    // Struct

    StructArray::StructArray(const std::shared_ptr<ArrayData> &data) {
        KCHECK_EQ(data->type->id(), Type::STRUCT);
        set_data(data);
        boxed_fields_.resize(data->child_data.size());
    }

    StructArray::StructArray(const std::shared_ptr<DataType> &type, int64_t length,
                             const std::vector<std::shared_ptr<Array>> &children,
                             std::shared_ptr<Buffer> null_bitmap, int64_t null_count,
                             int64_t offset) {
        KCHECK_EQ(type->id(), Type::STRUCT);
        set_data(ArrayData::create(type, length, {null_bitmap}, null_count, offset));
        for (const auto &child: children) {
            data_->child_data.push_back(child->data());
        }
        boxed_fields_.resize(children.size());
    }

    turbo::Result<std::shared_ptr<StructArray>> StructArray::create(
            const std::vector<std::shared_ptr<Array>> &children,
            const std::vector<std::shared_ptr<Field>> &fields,
            std::shared_ptr<Buffer> null_bitmap, int64_t null_count, int64_t offset) {
        if (children.size() != fields.size()) {
            return turbo::invalid_argument_error("Mismatching number of fields and child arrays");
        }
        if (children.empty()) {
            return turbo::invalid_argument_error("Can't infer struct array length with 0 child arrays");
        }
        const int64_t length = children.front()->length();
        for (const auto &child: children) {
            if (length != child->length()) {
                return turbo::invalid_argument_error("Mismatching child array lengths");
            }
        }
        if (offset > length) {
            return turbo::out_of_range_error("Offset greater than length of child arrays");
        }
        if (null_bitmap == nullptr) {
            if (null_count > 0) {
                return turbo::invalid_argument_error("null_count = ", null_count, " but no null bitmap given");
            }
            null_count = 0;
        }
        return std::make_shared<StructArray>(STRUCT(fields), length - offset, children,
                                             null_bitmap, null_count, offset);
    }

    turbo::Result<std::shared_ptr<StructArray>> StructArray::create(
            const std::vector<std::shared_ptr<Array>> &children,
            const std::vector<std::string> &field_names, std::shared_ptr<Buffer> null_bitmap,
            int64_t null_count, int64_t offset) {
        if (children.size() != field_names.size()) {
            return turbo::invalid_argument_error("Mismatching number of field names and child arrays");
        }
        std::vector<std::shared_ptr<Field>> fields(children.size());
        for (size_t i = 0; i < children.size(); ++i) {
            fields[i] = ::nebula::field(field_names[i], children[i]->type());
        }
        return create(children, fields, std::move(null_bitmap), null_count, offset);
    }

    const StructType *StructArray::struct_type() const {
        return turbo::checked_cast<const StructType *>(data_->type.get());
    }

    const ArrayVector &StructArray::fields() const {
        for (int i = 0; i < num_fields(); ++i) {
            (void) field(i);
        }
        return boxed_fields_;
    }

    const std::shared_ptr<Array> &StructArray::field(int i) const {
        std::shared_ptr<Array> result = std::atomic_load(&boxed_fields_[i]);
        if (!result) {
            std::shared_ptr<ArrayData> field_data;
            if (data_->offset != 0 || data_->child_data[i]->length != data_->length) {
                field_data = data_->child_data[i]->slice(data_->offset, data_->length);
            } else {
                field_data = data_->child_data[i];
            }
            std::shared_ptr<Array> result = make_array(field_data);
            std::atomic_store(&boxed_fields_[i], result);
            return boxed_fields_[i];
        }
        return boxed_fields_[i];
    }

    std::shared_ptr<Array> StructArray::get_field_by_name(const std::string &name) const {
        int i = struct_type()->get_field_index(name);
        return i == -1 ? nullptr : field(i);
    }

    turbo::Status StructArray::can_reference_field_by_name(const std::string &name) const {
        if (get_field_by_name(name) == nullptr) {
            return turbo::invalid_argument_error("Field named '", name,
                                                 "' not found or not unique in the struct.");
        }
        return turbo::OkStatus();
    }

    turbo::Status StructArray::can_reference_fields_by_names(
            const std::vector<std::string> &names) const {
        for (const auto &name: names) {
            TURBO_RETURN_NOT_OK(can_reference_field_by_name(name));
        }
        return turbo::OkStatus();
    }

    turbo::Result<ArrayVector> StructArray::Flatten(MemoryPool *pool) const {
        ArrayVector flattened;
        flattened.resize(data_->child_data.size());
        std::shared_ptr<Buffer> null_bitmap = data_->buffers[0];

        for (int i = 0; static_cast<size_t>(i) < data_->child_data.size(); i++) {
            TURBO_MOVE_OR_RAISE(flattened[i], GetFlattenedField(i, pool));
        }

        return flattened;
    }

    turbo::Result<std::shared_ptr<Array>> StructArray::GetFlattenedField(int index,
                                                                         MemoryPool *pool) const {
        std::shared_ptr<Buffer> null_bitmap = data_->buffers[0];

        auto child_data = data_->child_data[index]->copy();

        std::shared_ptr<Buffer> flattened_null_bitmap;
        int64_t flattened_null_count = kUnknownNullCount;

        // Need to adjust for parent offset
        if (data_->offset != 0 || data_->length != child_data->length) {
            child_data = child_data->slice(data_->offset, data_->length);
        }
        std::shared_ptr<Buffer> child_null_bitmap = child_data->buffers[0];
        const int64_t child_offset = child_data->offset;

        // The validity of a flattened datum is the logical AND of the struct
        // element's validity and the individual field element's validity.
        if (null_bitmap && child_null_bitmap) {
            TURBO_MOVE_OR_RAISE(
                    flattened_null_bitmap,
                    BitmapAnd(pool, child_null_bitmap->data(), child_offset, null_bitmap_data_,
                              data_->offset, data_->length, child_offset));
        } else if (child_null_bitmap) {
            flattened_null_bitmap = child_null_bitmap;
            flattened_null_count = child_data->null_count;
        } else if (null_bitmap) {
            if (child_offset == data_->offset) {
                flattened_null_bitmap = null_bitmap;
            } else {
                // If the child has an offset, need to synthesize a validity
                // buffer with an offset too
                TURBO_MOVE_OR_RAISE(flattened_null_bitmap,
                                    allocate_empty_bitmap(child_offset + data_->length, pool));
                CopyBitmap(null_bitmap_data_, data_->offset, data_->length,
                           flattened_null_bitmap->mutable_data(), child_offset);
            }
            flattened_null_count = data_->null_count;
        } else {
            flattened_null_count = 0;
        }

        auto flattened_data = child_data->copy();
        flattened_data->buffers[0] = flattened_null_bitmap;
        flattened_data->null_count = flattened_null_count;

        return make_array(flattened_data);
    }

    // ----------------------------------------------------------------------
    // UnionArray

    void UnionArray::set_data(std::shared_ptr<ArrayData> data) {
        this->Array::set_data(std::move(data));

        union_type_ = turbo::checked_cast<const UnionType *>(data_->type.get());

        KCHECK_GE(data_->buffers.size(), 2);
        raw_type_codes_ = data->GetValuesSafe<int8_t>(1, /*offset=*/0);
        boxed_fields_.resize(data_->child_data.size());
    }

    void SparseUnionArray::set_data(std::shared_ptr<ArrayData> data) {
        this->UnionArray::set_data(std::move(data));
        KCHECK_EQ(data_->type->id(), Type::SPARSE_UNION);
        KCHECK_EQ(data_->buffers.size(), 2);

        // No validity bitmap
        KCHECK_EQ(data_->buffers[0], nullptr);
    }

    void DenseUnionArray::set_data(const std::shared_ptr<ArrayData> &data) {
        this->UnionArray::set_data(data);

        KCHECK_EQ(data_->type->id(), Type::DENSE_UNION);
        KCHECK_EQ(data_->buffers.size(), 3);

        // No validity bitmap
        KCHECK_EQ(data_->buffers[0], nullptr);

        raw_value_offsets_ = data->GetValuesSafe<int32_t>(2, /*offset=*/0);
    }

    SparseUnionArray::SparseUnionArray(std::shared_ptr<ArrayData> data) {
        set_data(std::move(data));
    }

    SparseUnionArray::SparseUnionArray(std::shared_ptr<DataType> type, int64_t length,
                                       ArrayVector children,
                                       std::shared_ptr<Buffer> type_codes, int64_t offset) {
        auto internal_data = ArrayData::create(std::move(type), length,
                                               BufferVector{nullptr, std::move(type_codes)},
                /*null_count=*/0, offset);
        for (const auto &child: children) {
            internal_data->child_data.push_back(child->data());
        }
        set_data(std::move(internal_data));
    }

    turbo::Result<std::shared_ptr<Array>> SparseUnionArray::GetFlattenedField(
            int index, MemoryPool *pool) const {
        if (index < 0 || index >= num_fields()) {
            return turbo::invalid_argument_error("Index out of range: ", index);
        }
        auto child_data = data_->child_data[index]->copy();
        // Adjust the result offset/length to be absolute.
        if (data_->offset != 0 || data_->length != child_data->length) {
            child_data = child_data->slice(data_->offset, data_->length);
        }
        std::shared_ptr<Buffer> child_null_bitmap = child_data->buffers[0];
        const int64_t child_offset = child_data->offset;

        // Synthesize a null bitmap based on the union discriminant.
        // Make sure the bitmap has extra bits corresponding to the child offset.
        TURBO_MOVE_OR_RAISE(std::shared_ptr<Buffer> flattened_null_bitmap,
                            allocate_empty_bitmap(child_data->length + child_offset, pool));
        const int8_t type_code = union_type()->type_codes()[index];
        const int8_t *type_codes = raw_type_codes();
        int64_t offset = 0;
        internal::GenerateBitsUnrolled(flattened_null_bitmap->mutable_data(), child_offset,
                                       data_->length,
                                       [&] { return type_codes[offset++] == type_code; });

        // The validity of a flattened datum is the logical AND of the synthesized
        // null bitmap buffer and the individual field element's validity.
        if (child_null_bitmap) {
            BitmapAnd(flattened_null_bitmap->data(), child_offset, child_null_bitmap->data(),
                      child_offset, child_data->length, child_offset,
                      flattened_null_bitmap->mutable_data());
        }

        child_data->buffers[0] = std::move(flattened_null_bitmap);
        child_data->null_count = kUnknownNullCount;
        return make_array(child_data);
    }

    DenseUnionArray::DenseUnionArray(const std::shared_ptr<ArrayData> &data) {
        set_data(data);
    }

    DenseUnionArray::DenseUnionArray(std::shared_ptr<DataType> type, int64_t length,
                                     ArrayVector children, std::shared_ptr<Buffer> type_ids,
                                     std::shared_ptr<Buffer> value_offsets, int64_t offset) {
        auto internal_data = ArrayData::create(
                std::move(type), length,
                BufferVector{nullptr, std::move(type_ids), std::move(value_offsets)},
                /*null_count=*/0, offset);
        for (const auto &child: children) {
            internal_data->child_data.push_back(child->data());
        }
        set_data(internal_data);
    }

    turbo::Result<std::shared_ptr<Array>> DenseUnionArray::create(
            const Array &type_ids, const Array &value_offsets, ArrayVector children,
            std::vector<std::string> field_names, std::vector<type_code_t> type_codes) {
        if (value_offsets.type_id() != Type::INT32) {
            return turbo::failed_precondition_error("UnionArray offsets must be signed int32");
        }

        if (type_ids.type_id() != Type::INT8) {
            return turbo::failed_precondition_error("UnionArray type_ids must be signed int8");
        }

        if (type_ids.null_count() != 0) {
            return turbo::invalid_argument_error("Union type ids may not have nulls");
        }

        if (value_offsets.null_count() != 0) {
            return turbo::invalid_argument_error("Make does not allow nulls in value_offsets");
        }

        if (field_names.size() > 0 && field_names.size() != children.size()) {
            return turbo::invalid_argument_error("field_names must have the same length as children");
        }

        if (type_codes.size() > 0 && type_codes.size() != children.size()) {
            return turbo::invalid_argument_error("type_codes must have the same length as children");
        }

        BufferVector buffers = {nullptr, turbo::checked_cast<const Int8Array &>(type_ids).values(),
                                turbo::checked_cast<const Int32Array &>(value_offsets).values()};

        auto union_type = dense_union(children, std::move(field_names), std::move(type_codes));
        auto internal_data =
                ArrayData::create(std::move(union_type), type_ids.length(), std::move(buffers),
                        /*null_count=*/0, type_ids.offset());
        for (const auto &child: children) {
            internal_data->child_data.push_back(child->data());
        }
        return std::make_shared<DenseUnionArray>(std::move(internal_data));
    }

    turbo::Result<std::shared_ptr<Array>> SparseUnionArray::create(
            const Array &type_ids, ArrayVector children, std::vector<std::string> field_names,
            std::vector<int8_t> type_codes) {
        if (type_ids.type_id() != Type::INT8) {
            return turbo::failed_precondition_error("UnionArray type_ids must be signed int8");
        }

        if (type_ids.null_count() != 0) {
            return turbo::invalid_argument_error("Union type ids may not have nulls");
        }

        if (field_names.size() > 0 && field_names.size() != children.size()) {
            return turbo::invalid_argument_error("field_names must have the same length as children");
        }

        if (type_codes.size() > 0 && type_codes.size() != children.size()) {
            return turbo::invalid_argument_error("type_codes must have the same length as children");
        }

        BufferVector buffers = {nullptr, turbo::checked_cast<const Int8Array &>(type_ids).values()};
        auto union_type = sparse_union(children, std::move(field_names), std::move(type_codes));
        auto internal_data =
                ArrayData::create(std::move(union_type), type_ids.length(), std::move(buffers),
                        /*null_count=*/0, type_ids.offset());
        for (const auto &child: children) {
            internal_data->child_data.push_back(child->data());
            if (child->length() != type_ids.length()) {
                return turbo::invalid_argument_error(
                        "Sparse UnionArray must have len(child) == len(type_ids) for all children");
            }
        }
        return std::make_shared<SparseUnionArray>(std::move(internal_data));
    }

    std::shared_ptr<Array> UnionArray::field(int i) const {
        if (i < 0 ||
            static_cast<decltype(boxed_fields_)::size_type>(i) >= boxed_fields_.size()) {
            return nullptr;
        }
        std::shared_ptr<Array> result = std::atomic_load(&boxed_fields_[i]);
        if (!result) {
            std::shared_ptr<ArrayData> child_data = data_->child_data[i]->copy();
            if (mode() == UnionMode::SPARSE) {
                // Sparse union: need to adjust child if union is sliced
                // (for dense unions, the need to lookup through the offsets
                //  makes this unnecessary)
                if (data_->offset != 0 || child_data->length > data_->length) {
                    child_data = child_data->slice(data_->offset, data_->length);
                }
            }
            result = make_array(child_data);
            std::atomic_store(&boxed_fields_[i], result);
        }
        return result;
    }

}  // namespace nebula
